#[cfg(test)]
extern crate quickcheck;

/**
 * Implementation of iMEC steganography method
 * from Schroeder de Witt, Sokota, Zico Kolter, Foerster and Strohmeier, 2023
 */

use futures::{future::FutureExt};

use ordered_float::NotNan;

use tokio::sync::mpsc::UnboundedSender;

use crate::{DiscreteDistribution, SparseCoupling, mec::Entropic};

#[derive(Debug)]
pub enum DecodeError {
    NoNextCodeSymbol(),
    NoSuchToken(String),
}

pub enum EncodeError {
    // AllWeightsZero()
}

impl std::fmt::Display for DecodeError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl std::error::Error for DecodeError {}

pub mod markov_backend {
    tonic::include_proto!("markov_backend");
}

use markov_backend::continuation_client::ContinuationClient;

use self::markov_backend::{GetRequest, InitRequest, PutRequest};

// Logging for educational purposes
pub enum DataOfInterest {
    EncodeIteration(usize),
    FinalResult(String, String),
    EncodeResult(String),
    MuTable(Vec<(Vec<String>, Vec<f64>, f64)>, Option<usize>),
    GammaTable(Vec<(f64, (usize, usize))>),
    IStar(usize),
    EncodeStart(),
    DecodeStart(),
}

pub trait NextTokenDistributionProvider<W, V> {
    fn select(
        &mut self,
        next: W,
    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send;
    fn get(
        &mut self,
    ) -> impl std::future::Future<Output = Result<DiscreteDistribution<W>, anyhow::Error>> + Send;
    fn reset(
        &mut self,
        init: V,
    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send;
}

pub fn make_uniform<A>(a: usize, labels: Option<Vec<A>>) -> DiscreteDistribution<A> {
    DiscreteDistribution {
        d: vec![1.0 / (a as f64); a],
        labels,
    }
}

#[derive(Clone)]
pub struct GRPCMarkovChain {
    pub backend_client: ContinuationClient<tonic::transport::Channel>,
}

impl NextTokenDistributionProvider<String, String> for GRPCMarkovChain {
    fn select(
        &mut self,
        next: String,
    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send {
        let result = self
            .backend_client
            .set_next_word(PutRequest { token: next });
        std::future::IntoFuture::into_future(result).map(|_ok_response| Ok(()))
    }

    fn reset(
        &mut self,
        start: String,
    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send {
        let result = self.backend_client.reset(InitRequest { start });
        std::future::IntoFuture::into_future(result).map(|_ok_response| Ok(()))
    }

    fn get(
        &mut self,
    ) -> impl std::future::Future<Output = Result<DiscreteDistribution<String>, anyhow::Error>> + Send
    {
        let result = self.backend_client.get_probabilities(GetRequest {});
        std::future::IntoFuture::into_future(result).map(|response| {
            response
                .map(|resp| {
                    let to_inner = resp.into_inner();
                    let entries = to_inner.entries.iter().map(|pat| {
                        (
                            std::cmp::Reverse(NotNan::new(pat.prob).unwrap()),
                            pat.token.clone(),
                        )
                    });

                    let d_in = entries.clone().map(|x| *x.0.0);
                    let d_in_sum = d_in.clone().sum::<f64>();

                    DiscreteDistribution {
                        labels: Some(entries.clone().map(|x| x.1).collect()),
                        d: d_in.map(|x| x / d_in_sum).collect(),
                    }
                })
                .map_err(anyhow::Error::from)
        })
    }
}

impl<W> NextTokenDistributionProvider<W, ()> for DiscreteDistribution<W>
where
    W: Clone + Send,
{
    fn select(
        &mut self,
        _next: W,
    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send {
        // do nothing: this is a memoryless Markov chain
        Box::pin(async move { Ok(()) })
    }

    fn reset(
        &mut self,
        _start: (),
    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send {
        // do nothing: this is a memoryless Markov chain
        Box::pin(async move { Ok(()) })
    }

    fn get(
        &mut self,
    ) -> impl std::future::Future<Output = Result<DiscreteDistribution<W>, anyhow::Error>> + Send
    {
        Box::pin(async move { Ok(self.clone()) })
    }
}

pub fn compute_marginals(coup: &SparseCoupling) -> (Vec<f64>, Vec<f64>) {
    let imax = coup.iter().map(|x| x.1.0).max().unwrap();
    let jmax = coup.iter().map(|x| x.1.1).max().unwrap();
    let mut rowsums = vec![0.0; imax + 1];
    let mut colsums = vec![0.0; jmax + 1];

    coup.iter().for_each(|(m, (i, j))| {
        rowsums[*i] += m;
        colsums[*j] += m;
    });

    (rowsums, colsums)
}

pub fn pick_istar<A>(mu: &[DiscreteDistribution<A>]) -> usize {
    let istar = mu
        .iter()
        .map(|x| x.d.entropy())
        .enumerate()
        .max_by_key(|(_, x)| NotNan::new(*x).unwrap())
        .map(|(i, _)| i)
        .unwrap();
    istar
}

pub fn log_mu_table<A: std::fmt::Display>(
    tx: Option<&UnboundedSender<DataOfInterest>>,
    mu: &[DiscreteDistribution<A>],
) {
    let putative_istar = pick_istar(mu);

    if let Some(tx) = tx {
        tx.send(DataOfInterest::MuTable(
            mu.iter()
                .map(|d| {
                    (
                        d.labels
                            .as_ref()
                            .unwrap()
                            .iter()
                            .map(|l| format!("{}", l))
                            .collect(),
                        d.d.clone(),
                        d.d.entropy(),
                    )
                })
                .collect(),
            Some(putative_istar),
        ))
        .unwrap();
    }
}
